// Copyright 2022 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "components/services/screen_ai/proto/main_content_extractor_proto_convertor.h"

#include "base/check_op.h"
#include "base/containers/flat_set.h"
#include "base/strings/string_util.h"
#include "base/strings/stringprintf.h"
#include "components/services/screen_ai/proto/view_hierarchy.pb.h"
#include "ui/accessibility/ax_enum_util.h"
#include "ui/accessibility/ax_enums.mojom.h"
#include "ui/accessibility/ax_node.h"
#include "ui/accessibility/ax_tree.h"
#include "ui/gfx/geometry/rect_f.h"
#include "ui/gfx/geometry/size_f.h"

namespace {

// Converts a Chrome role to a Screen2x role as text.
// TODO(https://crbug.com/1341655): Remove if MainContentExtractor (Screen2x)
// training protos are generated directly by Chrome or Screen2x uses the same
// role texts. Screen2x role names are generated by |blink::AXObject::RoleName|
// and these two function should stay in sync.
std::string GetMainContentExtractorRoleFromChromeRole(ax::mojom::Role role) {
  std::string role_name = ui::ToString(role);

  static base::flat_set<ax::mojom::Role> roles_with_similar_name = {
      ax::mojom::Role::kAlert,       ax::mojom::Role::kArticle,
      ax::mojom::Role::kBanner,      ax::mojom::Role::kBlockquote,
      ax::mojom::Role::kButton,      ax::mojom::Role::kCaption,
      ax::mojom::Role::kCell,        ax::mojom::Role::kCode,
      ax::mojom::Role::kComment,     ax::mojom::Role::kComplementary,
      ax::mojom::Role::kDefinition,  ax::mojom::Role::kDialog,
      ax::mojom::Role::kDirectory,   ax::mojom::Role::kDocument,
      ax::mojom::Role::kEmphasis,    ax::mojom::Role::kFeed,
      ax::mojom::Role::kFigure,      ax::mojom::Role::kForm,
      ax::mojom::Role::kGrid,        ax::mojom::Role::kGroup,
      ax::mojom::Role::kHeading,     ax::mojom::Role::kLink,
      ax::mojom::Role::kList,        ax::mojom::Role::kLog,
      ax::mojom::Role::kMain,        ax::mojom::Role::kMarquee,
      ax::mojom::Role::kMath,        ax::mojom::Role::kMenu,
      ax::mojom::Role::kMark,        ax::mojom::Role::kMeter,
      ax::mojom::Role::kNavigation,  ax::mojom::Role::kNone,
      ax::mojom::Role::kNote,        ax::mojom::Role::kParagraph,
      ax::mojom::Role::kRegion,      ax::mojom::Role::kRow,
      ax::mojom::Role::kSearch,      ax::mojom::Role::kSlider,
      ax::mojom::Role::kStatus,      ax::mojom::Role::kStrong,
      ax::mojom::Role::kSubscript,   ax::mojom::Role::kSuggestion,
      ax::mojom::Role::kSuperscript, ax::mojom::Role::kSwitch,
      ax::mojom::Role::kTab,         ax::mojom::Role::kTable,
      ax::mojom::Role::kTerm,        ax::mojom::Role::kTime,
      ax::mojom::Role::kTimer,       ax::mojom::Role::kToolbar,
      ax::mojom::Role::kTooltip,     ax::mojom::Role::kTree,
  };
  if (roles_with_similar_name.find(role) != roles_with_similar_name.end())
    return role_name;

  static base::flat_set<ax::mojom::Role> roles_with_all_lowercase_name = {
      ax::mojom::Role::kAlertDialog,   ax::mojom::Role::kApplication,
      ax::mojom::Role::kCheckBox,      ax::mojom::Role::kColumnHeader,
      ax::mojom::Role::kContentInfo,   ax::mojom::Role::kListBox,
      ax::mojom::Role::kListItem,      ax::mojom::Role::kMenuBar,
      ax::mojom::Role::kMenuItem,      ax::mojom::Role::kMenuItemCheckBox,
      ax::mojom::Role::kMenuItemRadio, ax::mojom::Role::kRadioGroup,
      ax::mojom::Role::kRowGroup,      ax::mojom::Role::kRowHeader,
      ax::mojom::Role::kScrollBar,     ax::mojom::Role::kSearchBox,
      ax::mojom::Role::kSpinButton,    ax::mojom::Role::kTabList,
      ax::mojom::Role::kTabPanel,      ax::mojom::Role::kTreeItem,
  };
  if (roles_with_all_lowercase_name.find(role) !=
      roles_with_all_lowercase_name.end()) {
    return base::ToLowerASCII(role_name);
  }

  static base::flat_map<ax::mojom::Role, std::string>
      roles_with_different_name = {
          // Aria Roles.
          {ax::mojom::Role::kComboBoxGrouping, "combobox"},
          {ax::mojom::Role::kComboBoxSelect, "combobox"},
          {ax::mojom::Role::kContentDeletion, "deletion"},
          {ax::mojom::Role::kDocAbstract, "doc-abstract"},
          {ax::mojom::Role::kDocAcknowledgments, "doc-acknowledgments"},
          {ax::mojom::Role::kDocAfterword, "doc-afterword"},
          {ax::mojom::Role::kDocAppendix, "doc-appendix"},
          {ax::mojom::Role::kDocBackLink, "doc-backlink"},
          {ax::mojom::Role::kDocBiblioEntry, "doc-biblioentry"},
          {ax::mojom::Role::kDocBibliography, "doc-bibliography"},
          {ax::mojom::Role::kDocBiblioRef, "doc-biblioref"},
          {ax::mojom::Role::kDocChapter, "doc-chapter"},
          {ax::mojom::Role::kDocColophon, "doc-colophon"},
          {ax::mojom::Role::kDocConclusion, "doc-conclusion"},
          {ax::mojom::Role::kDocCover, "doc-cover"},
          {ax::mojom::Role::kDocCredit, "doc-credit"},
          {ax::mojom::Role::kDocCredits, "doc-credits"},
          {ax::mojom::Role::kDocDedication, "doc-dedication"},
          {ax::mojom::Role::kDocEndnote, "doc-endnote"},
          {ax::mojom::Role::kDocEndnotes, "doc-endnotes"},
          {ax::mojom::Role::kDocEpigraph, "doc-epigraph"},
          {ax::mojom::Role::kDocEpilogue, "doc-epilogue"},
          {ax::mojom::Role::kDocErrata, "doc-errata"},
          {ax::mojom::Role::kDocExample, "doc-example"},
          {ax::mojom::Role::kDocFootnote, "doc-footnote"},
          {ax::mojom::Role::kDocForeword, "doc-foreword"},
          {ax::mojom::Role::kDocGlossary, "doc-glossary"},
          {ax::mojom::Role::kDocGlossRef, "doc-glossref"},
          {ax::mojom::Role::kDocIndex, "doc-index"},
          {ax::mojom::Role::kDocIntroduction, "doc-introduction"},
          {ax::mojom::Role::kDocNoteRef, "doc-noteref"},
          {ax::mojom::Role::kDocNotice, "doc-notice"},
          {ax::mojom::Role::kDocPageBreak, "doc-pagebreak"},
          {ax::mojom::Role::kDocPageFooter, "doc-pagefooter"},
          {ax::mojom::Role::kDocPageHeader, "doc-pageheader"},
          {ax::mojom::Role::kDocPageList, "doc-pagelist"},
          {ax::mojom::Role::kDocPart, "doc-part"},
          {ax::mojom::Role::kDocPreface, "doc-preface"},
          {ax::mojom::Role::kDocPrologue, "doc-prologue"},
          {ax::mojom::Role::kDocPullquote, "doc-pullquote"},
          {ax::mojom::Role::kDocQna, "doc-qna"},
          {ax::mojom::Role::kDocSubtitle, "doc-subtitle"},
          {ax::mojom::Role::kDocTip, "doc-tip"},
          {ax::mojom::Role::kDocToc, "doc-toc"},
          {ax::mojom::Role::kGenericContainer, "generic"},
          {ax::mojom::Role::kGraphicsDocument, "graphics-document"},
          {ax::mojom::Role::kGraphicsObject, "graphics-object"},
          {ax::mojom::Role::kGraphicsSymbol, "graphics-symbol"},
          {ax::mojom::Role::kCell, "gridcell"},
          {ax::mojom::Role::kImage, "img"},
          {ax::mojom::Role::kContentInsertion, "insertion"},
          {ax::mojom::Role::kListBoxOption, "option"},
          {ax::mojom::Role::kProgressIndicator, "progressbar"},
          {ax::mojom::Role::kRadioButton, "radio"},
          {ax::mojom::Role::kSplitter, "separator"},
          {ax::mojom::Role::kTextField, "textbox"},
          {ax::mojom::Role::kTreeGrid, "treegrid"},
          // Reverse Roles
          {ax::mojom::Role::kHeader, "banner"},
          {ax::mojom::Role::kToggleButton, "button"},
          {ax::mojom::Role::kPopUpButton, "combobox"},
          {ax::mojom::Role::kFooter, "contentinfo"},
          {ax::mojom::Role::kMenuListOption, "menuitem"},
          {ax::mojom::Role::kComboBoxMenuButton, "combobox"},
          {ax::mojom::Role::kTextFieldWithComboBox, "combobox"}};

  const auto& item = roles_with_different_name.find(role);
  if (item != roles_with_different_name.end())
    return item->second;

  // Roles that are not in the above tree groups have uppercase first letter
  // names.
  role_name[0] = base::ToUpperASCII(role_name[0]);
  return role_name;
}

// TODO(https://crbug.com/1278249): Consider merging the following functions
// into a template, e.g. using std::is_same.
void AddAttribute(const std::string& name,
                  int value,
                  screenai::UiElement& ui_element) {
  screenai::UiElementAttribute attrib;
  attrib.set_name(name);
  attrib.set_int_value(value);
  ui_element.add_attributes()->Swap(&attrib);
}

void AddAttribute(const std::string& name,
                  const char* value,
                  screenai::UiElement& ui_element) {
  screenai::UiElementAttribute attrib;
  attrib.set_name(name);
  attrib.set_string_value(value);
  ui_element.add_attributes()->Swap(&attrib);
}

void AddAttribute(const std::string& name,
                  const std::string& value,
                  screenai::UiElement& ui_element) {
  screenai::UiElementAttribute attrib;
  attrib.set_name(name);
  attrib.set_string_value(value);
  ui_element.add_attributes()->Swap(&attrib);
}

// Creates the proto for |node|, setting its own and parent id respectively to
// |id| and |parent_id|. Updates |tree_dimensions| to include the bounds of the
// new node.
// Requires setting "child_ids" and "bounding_box" properties in next steps.
screenai::UiElement CreateUiElementProto(const ui::AXTree& tree,
                                         const ui::AXNode* node,
                                         int id,
                                         int parent_id,
                                         gfx::SizeF& tree_dimensions) {
  screenai::UiElement uie;

  const ui::AXNodeData& node_data = node->data();

  // ID.
  uie.set_id(id);

  // Attributes.
  // TODO(https://crbug.com/1278249): Get attribute strings from a Google3
  // export, also the experimental ones for the unittest.
  AddAttribute("axnode_id", static_cast<int>(node->id()), uie);
  const std::string& display_value =
      node_data.GetStringAttribute(ax::mojom::StringAttribute::kDisplay);
  if (!display_value.empty())
    AddAttribute("/extras/styles/display", display_value, uie);
  AddAttribute("/extras/styles/visibility",
               node_data.IsInvisible() ? "hidden" : "visible", uie);
  // Add extra CSS attributes, such as text-align, hierarchical level, font
  // size, and font weight supported by both AXTree/AXNode and screen2x.
  // Screen2x expects these properties to be in the string format, so we
  // convert them into string.
  int32_t int_attribute_value;
  if (node_data.GetIntAttribute(ax::mojom::IntAttribute::kTextAlign,
                                &int_attribute_value)) {
    AddAttribute("/extras/styles/text-align",
                 ui::ToString((ax::mojom::TextAlign)int_attribute_value), uie);
  }
  if (node_data.GetIntAttribute(ax::mojom::IntAttribute::kHierarchicalLevel,
                                &int_attribute_value)) {
    AddAttribute("hierarchical_level", int_attribute_value, uie);
  }
  // Get float attributes and store them as string attributes in the screenai
  // proto for the main content extractor (screen2x).
  float float_attribute_value;
  if (node_data.GetFloatAttribute(ax::mojom::FloatAttribute::kFontSize,
                                  &float_attribute_value)) {
    AddAttribute("/extras/styles/font-size",
                 base::StringPrintf("%.0fpx", float_attribute_value), uie);
  }
  if (node_data.GetFloatAttribute(ax::mojom::FloatAttribute::kFontWeight,
                                  &float_attribute_value)) {
    AddAttribute("/extras/styles/font-weight",
                 base::StringPrintf("%.0f", float_attribute_value), uie);
  }

  // This is a fixed constant for Chrome requests to Screen2x.
  AddAttribute("class_name", "chrome.unicorn", uie);
  AddAttribute("chrome_role",
               GetMainContentExtractorRoleFromChromeRole(node_data.role), uie);
  AddAttribute("text",
               node_data.GetStringAttribute(ax::mojom::StringAttribute::kName),
               uie);

  // Type and parent.
  uie.set_parent_id(parent_id);

  // Type.
  uie.set_type(node == tree.root() ? screenai::UiElementType::ROOT
                                   : screenai::UiElementType::VIEW);

  // Bounding Box.
  gfx::RectF bounds = tree.RelativeToTreeBounds(
      node, gfx::RectF(0, 0),
      /* offscreen= */ nullptr, /* clip_bounds= */ false,
      /* skip_container_offset= */ false);

  // Bounding Box Pixels.
  screenai::BoundingBoxPixels* bounding_box_pixels =
      new screenai::BoundingBoxPixels();
  bounding_box_pixels->set_top(bounds.y());
  bounding_box_pixels->set_left(bounds.x());
  bounding_box_pixels->set_bottom(bounds.bottom());
  bounding_box_pixels->set_right(bounds.right());
  uie.set_allocated_bounding_box_pixels(bounding_box_pixels);

  tree_dimensions.set_height(fmax(tree_dimensions.height(), bounds.bottom()));
  tree_dimensions.set_width(fmax(tree_dimensions.width(), bounds.right()));

  return uie;
}

// Adds the subtree of |node| to |proto| with pre-order traversal.
// Uses |next_unused_node_id| as the current node id and updates it for the
// children. Updates |tree_dimensions| to include the bounds of the new node.
void AddSubTree(const ui::AXTree& tree,
                const ui::AXNode* node,
                screenai::ViewHierarchy& proto,
                int& next_unused_node_id,
                int parent_id,
                gfx::SizeF& tree_dimensions) {
  // Ensure that node id and index are the same.
  DCHECK(proto.ui_elements_size() == next_unused_node_id);

  // Create and add proto.
  int current_node_id = next_unused_node_id;
  screenai::UiElement uie = CreateUiElementProto(tree, node, current_node_id,
                                                 parent_id, tree_dimensions);
  proto.add_ui_elements()->Swap(&uie);

  // Add children.
  std::vector<int> child_ids;
  for (auto it = node->AllChildrenBegin(); it != node->AllChildrenEnd(); ++it) {
    child_ids.push_back(++next_unused_node_id);
    AddSubTree(tree, it.get(), proto, next_unused_node_id, current_node_id,
               tree_dimensions);
  }

  // Add child ids.
  for (int child : child_ids)
    proto.mutable_ui_elements(current_node_id)->add_child_ids(child);
}

}  // namespace

namespace screen_ai {

std::string SnapshotToViewHierarchy(const ui::AXTreeUpdate& snapshot) {
  DCHECK(!snapshot.nodes.empty());

  // Deserialize the snapshot.
  ui::AXTree tree(snapshot);

  // To be computed based on the max dimensions of all elements in the tree.
  // TODO(https://crbug.com/1278249): Consider using combination of scroll
  // max and view port size to find the tree dimensions. Screen2x is getting the
  // size from the screenshot image of the tree.
  gfx::SizeF tree_dimensions;

  // Screen2x requires the nodes to come in PRE-ORDER, and have only positive
  // ids. |AddSubTree| traverses the |tree| in preorder and creates the
  // required proto.
  int next_unused_node_id = 0;
  screenai::ViewHierarchy proto;
  AddSubTree(tree, tree.root(), proto, next_unused_node_id, /*parent_id=*/-1,
             tree_dimensions);

  // If the tree has a zero dimension, there is nothing to send.
  if (tree_dimensions.IsEmpty())
    return "";

  // The bounds of the root item should be set to the snapshot size.
  proto.mutable_ui_elements(0)->mutable_bounding_box_pixels()->set_right(
      tree_dimensions.width());
  proto.mutable_ui_elements(0)->mutable_bounding_box_pixels()->set_bottom(
      tree_dimensions.height());
  DCHECK_EQ(proto.ui_elements(0).bounding_box().right(), 0);
  DCHECK_EQ(proto.ui_elements(0).bounding_box().top(), 0);

  // Set relative sizes.
  for (int i = 0; i < proto.ui_elements_size(); i++) {
    auto* bounding_box = proto.mutable_ui_elements(i)->mutable_bounding_box();
    const auto& bounding_box_pixels =
        proto.ui_elements(i).bounding_box_pixels();
    bounding_box->set_top(bounding_box_pixels.top() / tree_dimensions.height());
    bounding_box->set_left(bounding_box_pixels.left() /
                           tree_dimensions.width());
    bounding_box->set_bottom(bounding_box_pixels.bottom() /
                             tree_dimensions.height());
    bounding_box->set_right(bounding_box_pixels.right() /
                            tree_dimensions.width());
  }

  return proto.SerializeAsString();
}

const std::map<std::string, ax::mojom::Role>&
GetMainContentExtractorToChromeRoleConversionMapForTesting() {
  static std::map<std::string, ax::mojom::Role> contentExtractionToChromeRoles;

  if (contentExtractionToChromeRoles.empty()) {
    for (int i = static_cast<int>(ax::mojom::Role::kMinValue);
         i <= static_cast<int>(ax::mojom::Role::kMaxValue); i++) {
      auto role = static_cast<ax::mojom::Role>(i);
      contentExtractionToChromeRoles[GetMainContentExtractorRoleFromChromeRole(
          role)] = role;
    }
  }

  return contentExtractionToChromeRoles;
}

}  // namespace screen_ai
